from contextlib import AsyncExitStack
from itertools import chain
import sys
import time
from tkinter import Frame
from colorama.ansi import code_to_chars
import psutil
import pprint
from typing import *
from pathlib import Path
from queue import PriorityQueue
from collections import Counter

import tqdm
from colorama import Fore, Style


class HuffmanNode:
    def __init__(
        self,
        byte: Union[None, bytes] = None,
        freq: Union[None, int] = None,
        leaf: bool = False,
        left_node=None,
        right_node=None,
        parent=None,
        length: int = 1,
    ):
        self.bytes: Union[None, bytes] = byte
        self.freq: Union[None, int] = freq
        self.leaf: bool = leaf
        self.left_node = left_node
        self.right_node = right_node
        self.parent_node = parent
        self.length = length
        self.go_through_time = 0

    def __lt__(self, other):
        if self.bytes is not None and other.bytes is None:
            return True
        elif self.bytes is None and other.bytes is not None:
            return False
        elif self.bytes is None and other.bytes is None:
            return self.length < other.length
        else:
            return self.bytes > other.bytes

    def __repr__(self):
        return f"Node: byte={self.bytes}, freq={self.freq}"


class HuffmanTree:
    """
    Huffman Tree
    Attributes:
        table (Dict[bytes, str]): bin digits of bytes
    """

    def __init__(self, forest: Dict[bytes, int]):
        """
        construction method of Huffman Tree
        Args:
            forest (Dict[bytes, int]: bytes counting
        """
        self.forest: Dict[bytes, int] = forest
        self.queue = self._build_priority_queue(forest)
        self.tree: HuffmanNode = self._build_tree()
        self.table: Dict[Union[bytes, int], Union[int, bytes]] = self._get_code()

    def _build_tree(self) -> HuffmanNode:
        root = HuffmanNode()
        while not self.queue.empty():
            right_node = self.queue.get()[1]
            right_last = True
            if not self.queue.empty():
                left_node = self.queue.get()[1]
                right_last = False
                new_node = HuffmanNode(
                    byte=None,
                    freq=left_node.freq + right_node.freq,
                    leaf=False,
                    left_node=left_node,
                    right_node=right_node,
                    length=max(left_node.length, right_node.length) + 1,
                )
                left_node.parent_node = new_node
                right_node.parent_node = new_node
                self.queue.put((new_node.freq, new_node))
            else:
                root = right_node
        return root

    def _get_code(self) -> Dict[Union[bytes, int], Union[bytes, int, str]]:
        table: Dict[Union[bytes, int], Union[bytes, int, str]] = {}
        bin_num: List[str] = []
        node_ptr = self.tree
        traversal_stack: List[HuffmanNode] = [self.tree]
        while len(table) < len(self.forest):
            # 第一次经过并且左孩子非空
            if node_ptr.left_node is not None and node_ptr.go_through_time == 0:
                if node_ptr != traversal_stack[-1]:
                    traversal_stack.append(node_ptr)
                node_ptr.go_through_time += 1
                node_ptr = node_ptr.left_node
                bin_num.append("0")
            # 右孩子非空且没有经过
            elif node_ptr.right_node is not None and node_ptr.go_through_time <= 1:
                node_ptr.go_through_time += 1
                if node_ptr != traversal_stack[-1]:
                    traversal_stack.append(node_ptr)
                node_ptr = node_ptr.right_node
                bin_num.append("1")
            # 如果是叶节点
            elif node_ptr.leaf:
                # 不敢像下面这样直接从int转为1字节，8bit长度的字节，因为会补齐，导致压缩失去意义
                # table[node_ptr.bytes] = int("".join(bin_num), 2).to_bytes(length=1, byteorder="big", signed=False)
                # table[node_ptr.bytes] = bin(int("".join(bin_num), 2))
                table[node_ptr.bytes] = "".join(bin_num)
                bin_num.pop(-1)
                node_ptr = traversal_stack[-1]
            # 否则父节点已经经过了两次
            else:
                bin_num[-1] = "1"
                traversal_stack.pop(-1)
                node_ptr = traversal_stack[-1]
        return table

    @staticmethod
    def _build_priority_queue(forest: Dict[bytes, int]) -> PriorityQueue:
        queue = PriorityQueue()
        for byte, freq in (d := sorted(forest.items(), key=lambda x: x[1])) :
            # lower frequency, lower priority
            queue.put((freq, HuffmanNode(byte=byte, freq=freq, leaf=True, length=1)))
        if queue.empty():
            print(f"{Fore.GREEN}Empty file, pass{Style.RESET_ALL}")
            sys.exit(0)
        return queue


class HuffmanEncoding(object):
    def __init__(
        self, file_path: Union[str, Path], save_folder: Union[Path, None] = None
    ):
        self.file_path: Path = Path(file_path) if isinstance(
            file_path, str
        ) else file_path
        self.save_folder: Path = Path(
            save_folder
        ) if save_folder is not None else file_path.parent
        temp: Tuple[Dict[bytes, int], bytes]
        self.huffman_tree: HuffmanTree = HuffmanTree(
            (temp := self.read(self.file_path))[0]
        )
        self.raw_bytes: bytes = temp[1]
        self.encoded = self.encoding()

    def encoding(self) -> List[Union[int, bytes]]:
        return [self.huffman_tree.table[ii] for ii in self.raw_bytes]

    def save(self):
        pass

    @staticmethod
    def read(file_path: Path) -> Tuple[Dict[bytes, int], bytes]:
        # TODO: 大文件读取
        assert file_path.exists(), f"{Fore.RED}文件不存在！{Style.RESET_ALL}"
        # read file of any type as byte stream, each unit of 8 units, maximum 256
        with open(file_path, "rb") as f:
            raw_bytes = f.read()
        return Counter(raw_bytes), raw_bytes


if __name__ == "__main__":
    with open(
        "./D星演示.pptx",
        "rb",
    ) as f:
        c = f.read()
        import chardet

        print(chardet.detect(c))
    from dahuffman import HuffmanCodec

    print(c)
    code = HuffmanCodec.from_frequencies(dict(Counter(c)))
    code.print_code_table()
    with open("./test.huf", "wb") as f:
        f.write(enc := code.encode(c))
    print(enc)

    with open("./test.huf", "rb") as f:
        c = f.read()
    print(code.decode(c))
    with open("decode.pdf", "wb") as f:
        f.write(b"".join([i.to_bytes(1, byteorder="big", signed=False) for i in code.decode(c)]))
    code.print_code_table()
    # he = HuffmanEncoding(file_path=Path(__file__).resolve().parent.joinpath("./ss.txt"))
    # pprint.pprint(he.huffman_tree.table)
    # encoded = he.encoding()
    # # 不等长编码需要以等长的编码形式书写，关键在于凑够八个比特才进行一次转bit
    # a = "".join(encoded)
    # d = 0
    # with open("./test.huf", "wb") as f:
    #     for i in tqdm.tqdm(range(0, len(a), 8)):
    #         current_byte = a[i:i+8]
    #         if len(current_byte) < 8:
    #             current_byte = f"{current_byte:<08}"
    #         f.write((b:=int(current_byte, 2)).to_bytes(length=1, byteorder="big", signed=False))
    #         d += 1
    # print(d)
    # bs = []
    # # 大文件在这里会卡住，最后被kill
    # for i in range(0, len(a), 8):
    #     bs.append(a[i: i+8])
    # # 最后一位补0
    # if len(bs[-1]) < 8:
    #     bs[-1] = f"{bs[-1]:<08}"
    # pprint.pprint(bs)

    # with open("./test.huf", "wb") as f:
    #     for i in bs:
    #         f.write((a:=int(i, 2)).to_bytes(length=1, byteorder="big", signed=False))
    # print(encoded, type(encoded[0]), sep="\n")
    # print(a := he.encoding())
    # print(b := int(a, 2).to_bytes(length=len(a)//8, byteorder="big", signed=False))
